/*************************************************************************
 > File Name: monolith/interface/inference/ultra_face_dect_model.cc
 > Author: dzhang
 > Mail: dong.zhang@ingenic.com
 > Created Time : Sun 17 Jan 2021 10:11:17 PM CST
 > Description:
 ************************************************************************/

#include "monolith/interface/inference/ultra_face_dect_model.h"
#include "monolith/core/macros.h"
#include "monolith/interface/inference/utils.h"
#include "glog/logging.h"
CC_NAME_SPACE_BEGIN
namespace inference {

UltraFaceDectModel* UltraFaceDectModel::model_ = nullptr;

static UltraFaceDectModel* UltraFaceDectModel::get_instatnce(){
    CHECK(model_) << "Please init UltraFaceDectModel";
    return model_;
}

static int UltraFaceDectModel::init(ModelParameter model_parameter, float score_threshold =
        0.7, float iou_threshold = 0.3, int topk = -1){
    model_ = new UltraFaceDectModel(model_parameter, score_threshold, iou_threshold, topk);
    return 0;
}

UltraFaceDectModel::UltraFaceDectModel(ModelParameter model_parameter,
        float score_threshold, float iou_threshold, int topk) :
        Model(model_parameter) {
    score_threshold_ = score_threshold;
    iou_threshold_ = iou_threshold;
    w_h_list_ = {model_parameter.input_width_, model_parameter.input_height_};

    for (auto size : w_h_list_) {
        std::vector<float> fm_item;
        for (float stride : strides_) {
            fm_item.push_back(ceil(size / stride));
        }
        featuremap_size_.push_back(fm_item);
    }

    for (auto size : w_h_list_) {
        shrinkage_size_.push_back(strides_);
    }
    /* generate prior anchors */
    for (int index = 0; index < num_featuremap_; index++) {
        float scale_w = model_parameter.input_width_
                / shrinkage_size_[0][index];
        float scale_h = model_parameter.input_height_
                / shrinkage_size_[1][index];
        for (int j = 0; j < featuremap_size_[1][index]; j++) {
            for (int i = 0; i < featuremap_size_[0][index]; i++) {
                float x_center = (i + 0.5) / scale_w;
                float y_center = (j + 0.5) / scale_h;

                for (float k : min_boxes_[index]) {
                    float w = k / model_parameter.input_width_;
                    float h = k / model_parameter.input_height_;
                    priors_.push_back( { CLIP(x_center, 1), CLIP(y_center, 1),
                            CLIP(w, 1), CLIP(h, 1) });
                }
            }
        }
    }
    num_anchors_ = priors_.size();
}

int UltraFaceDectModel::run(uint8_t* data, std::vector<Box> &face_list,
        int input_width, int input_height, int input_channel) {
    CHECK(data) << "UltraFaceDectModel image is empty ,please check!";
    model_interpreter_->resizeTensor(input_tensor_, {1, 3, model_parameter_.input_height_, model_parameter_.input_width_});
    model_interpreter_->resizeSession(model_session_);
    std::shared_ptr<MNN::CV::ImageProcess> pretreat(
            MNN::CV::ImageProcess::create(MNN::CV::BGR, MNN::CV::RGB, model_parameter_.mean_.data(), 3,
                    model_parameter_.var_.data(), 3));
    pretreat->convert(data, model_parameter_.input_width_, model_parameter_.input_height_, model_parameter_.input_width_*model_parameter_.input_channel_, input_tensor_);
    model_interpreter_->runSession(model_session_);

    std::string scores = "scores";
    std::string boxes = "boxes";
    MNN::Tensor *tensor_scores = model_interpreter_->getSessionOutput(model_session_, scores.c_str());
    MNN::Tensor *tensor_boxes = model_interpreter_->getSessionOutput(model_session_, boxes.c_str());
    MNN::Tensor tensor_scores_host(tensor_scores, tensor_scores->getDimensionType());
    tensor_scores->copyToHostTensor(&tensor_scores_host);
    MNN::Tensor tensor_boxes_host(tensor_boxes, tensor_boxes->getDimensionType());
    tensor_boxes->copyToHostTensor(&tensor_boxes_host);
    std::vector<Box> bbox_collection;
    generateBBox(bbox_collection, tensor_scores, tensor_boxes, input_width, input_height);
    nms(bbox_collection, face_list);
    return 0;
}

void UltraFaceDectModel::generateBBox(std::vector<Box> &bbox_collection,
        MNN::Tensor *scores, MNN::Tensor *boxes, int input_width,
        int input_height) {
    for (int i = 0; i < num_anchors_; i++) {
        if (scores->host<float>()[i * 2 + 1] > score_threshold_) {
            Box rects;
            float x_center = boxes->host<float>()[i * 4] * center_variance_
                    * priors_[i][2] + priors_[i][0];
            float y_center = boxes->host<float>()[i * 4 + 1] * center_variance_
                    * priors_[i][3] + priors_[i][1];
            float w = exp(boxes->host<float>()[i * 4 + 2] * size_variance_)
                    * priors_[i][2];
            float h = exp(boxes->host<float>()[i * 4 + 3] * size_variance_)
                    * priors_[i][3];

            rects.lx = CLIP(x_center - w / 2.0, 1) * input_width;
            rects.ly = CLIP(y_center - h / 2.0, 1) * input_height;
            rects.rx = CLIP(x_center + w / 2.0, 1) * input_width;
            rects.ry = CLIP(y_center + h / 2.0, 1) * input_height;
            rects.score = CLIP(scores->host<float>()[i * 2 + 1], 1);
            bbox_collection.push_back(rects);
        }
    }
}

} // namespace inference
CC_NAME_SPACE_END
